import time
from typing import List, Dict
from .memoryitem import MemoryItem
from sklearn.metrics.pairwise import cosine_similarity
from memory.keywordextractor import KeywordExtractor
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

class ImportanceCalculator:
    def __init__(self):
        self.vectorizer = TfidfVectorizer() 
        self.time_decay_rate = 0.8  
        self.default_intent_words = ["重要", "记住"]

    def calculate(
        self,
        current_text: str,
        context: List[str],
        long_term_items: List[MemoryItem],
        user_intent_words: List[str] = None,
        message_timestamp: float = None
    ) -> float:
        """计算重要性得分（0~1范围）"""
        keywords = KeywordExtractor().extract(current_text, context)
        keyword_score = 0
        if keywords:
            total_matches = sum(1 for item in long_term_items for kw in keywords if kw in item.text)
            keyword_score = total_matches / len(keywords) if keywords else 0

        if not context:
            semantic_score = 0
        else:
            vectors = self.vectorizer.fit_transform(context + [current_text])
            text_vec = vectors[-1]          
            context_vecs = vectors[:-1]    
            semantic_score = cosine_similarity(text_vec, context_vecs).mean()

        intent_words = user_intent_words or self.default_intent_words
        intent_score = 0.3 if any(word in current_text for word in intent_words) else 0

        if message_timestamp is None:
            message_timestamp = time.time()
        time_decay = self.time_decay_rate ** ((time.time() - message_timestamp) / 86400)

        weights = {
            "keyword": 0.4,
            "semantic": 0.3,
            "intent": 0.2,
            "time": 0.1
        }
        total_score = (
            weights["keyword"] * keyword_score +
            weights["semantic"] * semantic_score +
            weights["intent"] * intent_score +
            weights["time"] * time_decay
        )
        return min(max(total_score, 0), 1) 
